
__author__ = "Adrian Price-Whelan <adrn@nyu.edu>"

import os
import pyfits as pf
import numpy as np

# Color the output
_okColor = '\033[92m'
_errorColor = '\033[91m'
_ENDC = '\033[0m'


def hist(r,g,b):
    import matplotlib.pyplot as plt
    plt.subplot(311)
    plt.hist(r.ravel(), bins=175)
    plt.subplot(312)
    plt.hist(g.ravel(), bins=175)
    plt.subplot(313)
    plt.hist(b.ravel(), bins=175)
    plt.show()
    sys.exit(0)
    
def fitsToBands(imPath, rggb=False, bayer=True):
    """ 
    This function takes a path to a FITS file as an input, and decomposes 
    the image into separate bands. This function assumes the image is in
    the Primary HDU (HDU 0) of the specified file, and makes some 
    assumptions about what to return depending on the number of axes in 
    the image. If NAXIS=1, the 1D array is simply returned. If NAXIS=2, 
    it by default assumes a Bayer pattern, but this can be controlled by
    the 'bayer' parameter (True or False). If Bayer is False, the full 2D
    array is returned. If 'bayer' is True, and 'rggb' is False (default), 
    the function combines G1 and G2. If 'rggb' is True, G1 and G2 are 
    returned separately. If NAXIS=3, return a 2D array for each element in
    NAXIS3. If NAXIS > 3, an error is thrown. 

    Parameters
    ----------
    imPath : string
        A string representing a path to a FITS file.
    rggb : Boolean
        If NAXIS=2 and the image is stored in a Bayer pattern, this 
        parameter controls whether G1 and G2 are averaged or not.
    bayer : Boolean
        If NAXIS=2, this parameter controls whether the full array is
        returned, or if it is decomposed into individual bands.

    Examples
    --------
    >>> 

    """
    
    if not os.path.exists(imPath): raise IOError("Image: %s not found! Is the path correct? \n\n" % imPath)
    
    im = pf.open(imPath)
    naxis = im[0].header['NAXIS']
    imData = im[0].data
    im.close()
    
    if naxis == 1:
        return imData.astype(float)
        
    elif naxis == 2:
        if bayer:
            # The image is saved in a Bayer pattern
            R = imData[::2,::2].astype(float)
            G1 = imData[1::2,::2].astype(float)
            G2 = imData[::2,1::2].astype(float)
            B = imData[1::2,1::2].astype(float)
            
            # Sometimes CCD's don't have an even number of pixels (WTF??),
            #   e.g. the (old?) Nikon D40 X. This corrects for any shape
            #   mismatching by throwing out a row or column so that all 
            #   of the bands have the same shape
            if (R.shape[0] % 2) != 0: R = R[:-1,:]
            if (R.shape[1] % 2) != 0: R = R[:,:-1]
            
            if (G1.shape[0] % 2) != 0: G1 = G1[:-1,:]
            if (G1.shape[1] % 2) != 0: G1 = G1[:,:-1]
            
            if (G2.shape[0] % 2) != 0: G2 = G2[:-1,:]
            if (G2.shape[1] % 2) != 0: G2 = G2[:,:-1]
                
            if (B.shape[0] % 2) != 0: B = B[:-1,:]
            if (B.shape[1] % 2) != 0: B = B[:,:-1]
            
            if rggb:
                return (R, G1, G2, B)
            else:
                G = (G1 + G2) / 2.0
                return np.array([R, G, B])
        else:
            # The image is grayscale
            return imData.astype(float)
            
    elif naxis == 3:
        # NAXIS3 appears as the **first** index in the returned data
        return np.array([oneBand.astype(float) for oneBand in imData])

def luptonScale(r, g, b, beta=3.0, scales=[1.,1.,1.], axis=0):
    """ 
    This is an implementation of Lupton et. al. 2004's arcsinh scaling
    function for preserving color in astronomical images. Many RGB images
    effectively lose information, as most stars are saturated and end up
    white. See: http://arxiv.com/pdf/astro-ph/0312483 for more detail.

    Parameters
    ----------
    r : numpy.ndarray
        A 2D array to be mapped to the g-band in the output image.
    g : numpy.ndarray
        A 2D array to be mapped to the g-band in the output image.
    b : numpy.ndarray
        A 2D array to be mapped to the g-band in the output image.
    beta : float or any numeric type
        The so-called "softening parameter" - controls the non-linearity
        of the arcsinh stretch.
    scales : list
        Scale the bands by the specified amounts.

    """
    
    r = np.asarray(r, dtype=float)
    g = np.asarray(g, dtype=float)
    b = np.asarray(b, dtype=float)
    arr = np.array([r, g, b])
    
    imShape = r.shape
    if g.shape != imShape or b.shape != imShape:
        raise ValueError("The 3 bands must have the same shape!")
    
    radiusI = np.sum(arr, axis=axis)/3.0 * beta
    fac = np.where(radiusI > 0, np.arcsinh(radiusI)/radiusI, 0.0)
    
    if axis == 0: return np.array([arr[ii] * fac * scales[ii] for ii in range(arr.shape[0])])
    elif axis == 1: return np.array([arr[:,ii] * fac * scales[ii] for ii in range(arr.shape[1])])
    elif axis == 2: return np.array([arr[:,:,ii] * fac * scales[ii] for ii in range(arr.shape[2])])
    else: raise ValueError("Axis can only be 0, 1, or 2")

def normalize(arr):
    nArr = np.asarray(arr)
    return (nArr - nArr.min()) / (nArr.max() - nArr.min())

def normalizeRGB(r,g,b):
    return normalize(r), normalize(g), normalize(b)

def skySubtract(arr):
    # From Andy Becker..
    arrMode = 2.5*np.median(arr) - 1.5*np.mean(arr)
    return arr-arrMode